from __future__ import with_statement
from khronos.utils import INF, Logger, CPUTracker, TimeDelta, Namespace

indline_fmt = "%40s = %s"

class MultiStart(object):
    def __init__(self, log=True):
        self.log = Logger(enabled=log)
        self.instance = None
        self.upper_bound = INF
        self.optimal = -INF
        self.solutions = []
        self.meta = Namespace(Iterations=0)
        self.running = False
        self.cpu = CPUTracker()
        self.cpu_limit = INF
        self.iteration_limit = INF
        self.improvement_limit = INF
        
    def solve(self, instance, cpu=INF, iterations=INF, improvements=INF, 
              upper_bound=INF, optimal=-INF):
        """Initialize and run in a single method."""
        self.init(instance, upper_bound, optimal)
        self.run(cpu, iterations, improvements)
        
    def init(self, instance=None, upper_bound=INF, optimal=-INF):
        """Initialize the multi-start procedure."""
        if instance is not None:
            self.instance = instance
        self.upper_bound = upper_bound
        self.optimal = optimal
        self.solutions = []
        self.metadata = Namespace(Metaheuristic=self.__class__.__name__, 
                                  Instance=self.instance, 
                                  Iterations=0)
        self.running = False
        self.cpu = CPUTracker()
        self.cpu_limit = INF
        self.iteration_limit = INF
        self.improvement_limit = INF
        self.log(" --- Multi-start initialized --- ")
        
    def step(self, iterations=1):
        """Shortcut method. Equivalent to run(iterations=xxx)."""
        self.run(iterations=iterations)
        
    def improve(self, improvements=1):
        """Shortcut method. Equivalent to run(improvements=xxx)."""
        self.run(improvements=improvements)
        
    def run(self, cpu=INF, iterations=INF, improvements=INF):
        """The main loop of the multi-start 'metaheuristic'. Its behavior consists of simply 
        repeating the following steps:
            1) Build a solution.
            2) Check if the solution's z-value is better than the current upper bound. If so, 
                then update the upper bound and record the solution.
            3) Go to step 1."""
        self.cpu_limit = self.cpu.total() + cpu
        self.iteration_limit = self.meta.Iterations + iterations
        self.improvement_limit = len(self.solutions) + improvements
        with self.cpu.tracking():
            self.running = True
            try:
                while self.running:
                    solution = self.build_solution()
                    z = self.objective(solution)
                    if z < self.upper_bound:
                        self.log("IMPROVEMENT: %s -> %s" % (self.upper_bound, z))
                        self.add_solution(solution, z)
                        self.upper_bound = z
                        if self.upper_bound <= self.optimal:
                            self.log("OPTIMAL SOLUTION: stopping search.")
                            self.running = False
                        elif len(self.solutions) >= self.improvement_limit:
                            self.log("IMPROVEMENT LIMIT REACHED: stopping search.")
                            self.running = False
                            
                    self.meta.Iterations += 1
                    if self.meta.Iterations >= self.iteration_limit:
                        self.log("ITERATION LIMIT: %d iterations completed." % 
                                 (self.meta.Iterations,))
                        self.running = False
                    elif self.cpu.total() >= self.cpu_limit:
                        self.log("CPU LIMIT: %s cpu seconds elapsed." % (self.cpu.total(),))
                        self.running = False
            except KeyboardInterrupt:
                self.log("INTERRUPT: search interrupted manually.")
                self.running = False
                
    def add_solution(self, solution, z):
        """Add an improving solution to the solution list. The solution is coupled with a 
        metadata dictionary containing a snapshot of the current status of the solver."""
        meta = Namespace(self.meta)
        meta.update(Objective=z, Solution=solution, 
                    CPU=TimeDelta(seconds=self.cpu.total()), 
                    Search_Speed=(self.meta.Iterations + 1) / self.cpu.total())
        self.solutions.append(meta)
        # Log the solution's metadata.
        self.log("=" * 80)
        for key, value in sorted(meta.iteritems()):
            self.log(indline_fmt % (key, value), width=80)
        self.log("=" * 80)
        
    def best_solution(self):
        """If at least one solution has been found, return the best of them. Note that the list 
        of solutions is ordered by decreasing z-value, so the last element of the list contains 
        our best solution. If the list is empty, None is returned.
        NOTE: a solution is a tuple consisting of (node, metadata), where the z-value (along 
        with other information) is placed inside the metadata dictionary."""
        if len(self.solutions) > 0:
            return self.solutions[-1]
        return None
        
    # --------------------------------------------
    def report(self):
        """Print a summary of the solver's configuration, current search state and a summary 
        (z-value and elapsed cpu time) of the solutions found so far."""
        print object.__repr__(self)
        if self.instance is not None:
            gss = "%.3f iter/sec" % (self.meta.Iterations / self.cpu.total(),)
            print indline_fmt % ("Global search speed (GSS)", gss)
            print indline_fmt % ("Total CPU time", TimeDelta(seconds=self.cpu.total()))
            print indline_fmt % ("Iterations completed", self.meta.Iterations)
            print indline_fmt % ("Upper bound", self.upper_bound)
            print indline_fmt % ("Solutions found", len(self.solutions))
            prev_iter = 0
            prev_cpu = 0.0
            for x, sol in enumerate(self.solutions):
                iss = (sol.Iterations + 1 - prev_iter) / (sol.CPU.to_seconds() - prev_cpu)
                prev_iter = sol.Iterations + 1
                prev_cpu = sol.CPU.to_seconds()
                sol_info = "%s (obtained at %s, ISS=%.3f iter/sec)" % \
                    (sol.Objective, sol.CPU, iss)
                print indline_fmt % ("Solution %d" % (x,), sol_info)
                
    def plot_solutions(self, filename=None, dpi=300):
        """Build and display or save a plot showing the evolution of the upper bound (aka the 
        best solution) over time. If filename is None, the plot is displayed to the screen, 
        otherwise it is saved to the specified file. The filename extension will determine the 
        format of the output file (e.g. pdf, eps, png, jpg)."""
        if len(self.solutions) == 0:
            print "Error: No solutions to be plotted."
            return
        meta = self.solutions[0]
        cpu = meta.CPU.to_seconds()
        objective = meta.Objective
        xs = [cpu]
        ys = [objective]
        for meta in self.solutions[1:]:
            cpu = meta.CPU.to_seconds()
            xs.append(cpu)
            ys.append(objective)
            objective = meta.Objective
            xs.append(cpu)
            ys.append(objective)
        xs.append(self.cpu.total())
        ys.append(self.upper_bound)
        
        from matplotlib.pyplot import figure
        fig = figure()
        plot = fig.add_subplot(1, 1, 1)
        plot.set_title("Upper bound over time")
        plot.set_xlabel("CPU time (s)")
        plot.set_ylabel("Upper bound")
        plot.plot(xs, ys, color="blue", linewidth=2, 
                  marker="o", markeredgecolor="blue")
        if filename is None:
            fig.show()
        else:
            fig.savefig(filename, dpi=dpi)
            
    # --------------------------------------------------------------------------
    def build_solution(self):
        """This method should be redefined in subclasses to provide solutions to the
        multi-start main loop."""
        raise NotImplementedError()
        
    def objective(self, solution):
        """Compute and return the objective function value of the given solution."""
        raise NotImplementedError()
        
